import numpy as np
import matplotlib.pyplot as plt
from init_control import get_control_parameters
from pid_controller import PIDController, CascadePIDController
from uav_dynamics import uav_dynamics
from scipy.spatial.transform import Rotation as R

def motor_mixer(roll, pitch, yaw, thrust):
    """
    Control allocation. The quadrotor type is X-configuration,
    and the airframe is as follows:
    3↓   1↑
      \ /
      / \
    2↑   4↓

    Input:
      roll, pitch, yaw: attitude controller output.
      thrust: total thrust.

    Output:
      M1, M2, M3, M4: motor speed commands.
    """
    idle_PWM = 1000
    scale = 1000

    M1 = (thrust - roll + pitch + yaw) * scale + idle_PWM
    M2 = (thrust + roll - pitch + yaw) * scale + idle_PWM
    M3 = (thrust + roll + pitch - yaw) * scale + idle_PWM
    M4 = (thrust - roll - pitch - yaw) * scale + idle_PWM

    return M1, M2, M3, M4
def desired_euler_angles(eax, eay, psi):
    """
    Convert the desired horizontal acceleration to the desired Euler angle

    Input:
      eax, eay: desired horizontal acceleration

    Output:
      phi, theta: desired roll and pitch
    """
    g = 9.8  # 重力加速度 (m/s^2)

    # phi = eay / g
    # theta = -eax / g
    phi = (-np.sin(psi) * eax + np.cos(psi) * eay) / g
    theta = (-np.cos(psi) * eax - np.sin(psi) * eay) / g

    return phi, theta
def pwm_to_rpm(pwm):
    """
    Convert PWM signal to RPM. This is a simple linear approximation.
    You might need to adjust the conversion based on actual hardware data.

    Input:
      pwm: PWM signal value

    Output:
      rpm: motor speed in RPM
    """
    min_pwm = 1000
    max_pwm = 2000
    min_rpm = 0
    max_rpm = 10000  # 假设电机最大转速为10000 RPM

    # 限制PWM范围
    pwm = np.clip(pwm, min_pwm, max_pwm)

    # 线性插值计算RPM
    rpm = np.clip((pwm - min_pwm) / (max_pwm - min_pwm) ,-1,1)
    return rpm

def run_uav_simulation():
    params = get_control_parameters()
    DEG2RAD = params['DEG2RAD']
    x_position_controller = PIDController(params['Kpxp'], 0, 0, integrator_max=1,
                                          integrator_min=-1,max=5,min=-5)
    x_velocity_controller = PIDController(params['Kvxp'], params['Kvxi'], params['Kvxd'], integrator_max=1,
                                               integrator_min=-1,min=-0.4,max=0.4)
    y_position_controller = PIDController(params['Kpyp'], 0, 0, integrator_max=1,
                                          integrator_min=-1,max=5,min=-5,)
    y_velocity_controller = PIDController(params['Kvyp'], params['Kvyi'], params['Kvyd'], integrator_max=1,
                                          integrator_min=-1,min=-0.4,max=0.4)
    roll_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0, integrator_max=1,
                                          integrator_min=-1,max = DEG2RAD*35,min = -DEG2RAD*35)
    roll_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
                                         params['Kd_RP_AngleRate'], integrator_max=1,
                                         integrator_min=-1,max = DEG2RAD*220,min = -DEG2RAD*220)
    roll_controller = CascadePIDController(roll_angle_controller, roll_rate_controller)

    pitch_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0, integrator_max=1,
                                           integrator_min=-1,max = DEG2RAD*35,min = -DEG2RAD*35)
    pitch_rate_controller = PIDController(params['Kp_RP_AngleRate'], params['Ki_RP_AngleRate'],
                                          params['Kd_RP_AngleRate'], integrator_max=1,
                                          integrator_min=-1,max = DEG2RAD*220,min = -DEG2RAD*220)
    pitch_controller = CascadePIDController(pitch_angle_controller, pitch_rate_controller)
    yaw_angle_controller = PIDController(params['Kp_RP_ANGLE'], 0, 0, integrator_max=1,
                                           integrator_min=-1,max = DEG2RAD*35,min = -DEG2RAD*35)
    yaw_rate_controller = PIDController(params['Kp_YAW_AngleRate'], params['Ki_YAW_AngleRate'],
                                        params['Kd_YAW_AngleRate'], integrator_max=1,
                                        integrator_min=-1,max = DEG2RAD*220,min = -DEG2RAD*220)
    yaw_controller = CascadePIDController(yaw_angle_controller, yaw_rate_controller)

    height_position_controller = PIDController(params['Kpzp'], 0, 0, integrator_max=1,
                                               integrator_min=-1,min=-3,max=3)
    height_velocity_controller = PIDController(params['Kvzp'], params['Kvzi'], params['Kvzd'], integrator_max=1,
                                               integrator_min=-1,min=-0.4,max=0.4)
    x_controller = CascadePIDController(x_position_controller, x_velocity_controller)
    y_controller = CascadePIDController(y_position_controller, y_velocity_controller)
    height_controller = CascadePIDController(height_position_controller, height_velocity_controller)

    # 初始化仿真状态
    position = params['ModelInit_PosE'].astype(np.float64)
    velocity_world = np.array([0, 0, 0], dtype=np.float64)
    velocity_body = np.array([0, 0, 0], dtype=np.float64)
    angle = params['ModelInit_AngEuler'].astype(np.float64)
    rate = params['ModelInit_RateB'].astype(np.float64)
    # w_current = np.zeros(4, dtype=np.float64)  # 电机转速
    w_current = params['ModelInit_Rads']
    mass = params['ModelParam_uavMass']
    J = params['ModelParam_uavJ']
    ModelParam_motorCr = params['ModelParam_motorCr']
    ModelParam_motorWb = params['ModelParam_motorWb']

    dt = 0.011  # 时间步长，单位为秒
    simulation_time = 30  # 仿真总时间，单位为秒
    time_steps = int(simulation_time / dt)

    # 定义起飞阶段和正方形轨迹的四个角点

    target_altitude = -5  # 目标高度
    takeoff_time = 0  # 起飞时间，单位为秒
    takeoff_steps = int(takeoff_time / dt)

    square_corners = [
        params['ModelInit_PosE'].astype(np.float64),
        np.array([0, 0, target_altitude]),  # 起点
        np.array([5, 0, target_altitude]),
        np.array([5, 5, target_altitude]),
        np.array([0, 5, target_altitude]),
        np.array([0, 0, target_altitude]), # 回到起点
        np.array([0, 0, 0])  # 回到起点
    ]

    current_corner_index = 0
    next_corner_time = (time_steps - takeoff_steps) // len(square_corners)

    # 数据记录
    position_history = [position]
    setpoint_history = []

    for step in range(time_steps):
        if step < takeoff_steps:
            # 起飞阶段保持在初始点
            setpoint = np.array([0, 0, target_altitude * ((step) / takeoff_steps)])
        else:
            # 计算当前段的起点和终点
            start_corner = square_corners[current_corner_index]
            end_corner = square_corners[(current_corner_index + 1) % len(square_corners)]

            # 计算在当前段中的位置
            segment_duration = next_corner_time
            segment_step = (step - takeoff_steps) % segment_duration
            t = segment_step / segment_duration  # 归一化时间

            setpoint = (1 - t) * start_corner + t * end_corner
        print('setpoiont:',setpoint)
        # 位置控制器
        # eax = (setpoint[0] - position[0]) / (time_steps - takeoff_steps)
        # eay = (setpoint[1] - position[1]) / (time_steps - takeoff_steps)
        eax = x_controller.update(outer_setpoint=setpoint[0], outer_measurement=position[0],
                                                   inner_measurement=velocity_world[0], dt=dt)
        eay = y_controller.update(outer_setpoint=setpoint[1], outer_measurement=position[1],
                                  inner_measurement=velocity_world[1], dt=dt)
        print('eax,eay:',eax,eay)
        # phi, theta = desired_euler_angles(eax, eay,angle[2])
        phi, theta = desired_euler_angles(eax, eay, angle[2])
        phi = np.clip(phi,-220*DEG2RAD,220*DEG2RAD)
        theta = np.clip(theta, -220 * DEG2RAD, 220 * DEG2RAD)
        print('phi,theta:',phi,theta)
        roll_output = np.clip(roll_controller.update(outer_setpoint=phi, outer_measurement=angle[0],
                                             inner_measurement=rate[0], dt=dt),-1,1)
        pitch_output = np.clip(pitch_controller.update(outer_setpoint=theta, outer_measurement=angle[1],
                                               inner_measurement=rate[1], dt=dt),-1,1)
        # yaw_output = yaw_rate_controller.update(setpoint=0, measurement=angle[2], dt=dt)
        yaw_output = np.clip(yaw_controller.update(outer_setpoint=0, outer_measurement=angle[2],
                                               inner_measurement=rate[2], dt=dt),-1,1)
        altitude_output = height_controller.update(outer_setpoint=setpoint[2], outer_measurement=position[2],
                                                   inner_measurement=velocity_world[2], dt=dt)
        thrust = np.clip(-altitude_output + params['THR_HOVER'], 0.05, 0.9)
        print('roll_output,pitch_output,yaw_output,altitude_output:',roll_output,pitch_output,yaw_output,altitude_output)
        M1, M2, M3, M4 = motor_mixer(roll_output, pitch_output, yaw_output, thrust)
        print('M1,M2,M3,M4:',M1,M2,M3,M4)

        w_target = (np.array([pwm_to_rpm(M1), pwm_to_rpm(M2), pwm_to_rpm(M3), pwm_to_rpm(M4)])*ModelParam_motorCr+
                    np.array([ModelParam_motorWb,ModelParam_motorWb,ModelParam_motorWb,ModelParam_motorWb]))
        print('w_target:',w_target)
        dcm = dcm_from_euler(angle)
        total_force, total_torque,w_current = uav_dynamics(position, velocity_body, dcm, rate, w_current,w_target, dt, params)

        R = rotation_matrix(angle)
        acceleration_body = np.clip(total_force / mass,-0.4,0.4)
        # acceleration_world = R @ acceleration_body
        # print('acc:',acceleration_world)
        velocity_body+=acceleration_body*dt
        velocity_body = np.clip(velocity_body,-3,3)
        # velocity_world = np.clip(velocity_world+acceleration_world * dt,-1.3,1.3)
        velocity_world = R @ velocity_body
        position += velocity_world * dt
        position_history.append(position.copy())
        setpoint_history.append(setpoint.copy())
        print(f"Time: {step * dt:.2f}s, Position: {position}, Velocity_world: {velocity_world}, Angle: {angle}, Rate: {rate}")
        angular_acceleration = np.linalg.inv(J).dot(total_torque - np.cross(rate, J @ rate))
        rate += angular_acceleration * dt
        angle += rate * dt

        if step >= takeoff_steps and (step - takeoff_steps) % next_corner_time == 0 and step != takeoff_steps:
            current_corner_index = (current_corner_index + 1) % len(square_corners)

    position_history = np.array(position_history)
    setpoint_history = np.array(setpoint_history)

    fig = plt.figure()
    ax = fig.add_subplot(121, projection='3d')
    ax.plot(position_history[:, 0], position_history[:, 1], -position_history[:, 2], label='UAV Trajectory')
    ax.set_xlabel('X Position (m)')
    ax.set_ylabel('Y Position (m)')
    ax.set_zlabel('Z Position (m)')
    ax.set_title('Trajectory of UAV ')
    ax.legend()
    # plt.show()
    # fig = plt.figure()
    ax = fig.add_subplot(122, projection='3d')
    ax.plot(setpoint_history[:, 0], setpoint_history[:, 1], -setpoint_history[:, 2], label='setpoint Trajectory')
    ax.set_xlabel('X Position (m)')
    ax.set_ylabel('Y Position (m)')
    ax.set_zlabel('Z Position (m)')
    ax.set_title(' Desired Trajectory ')
    ax.legend()
    plt.show()

def rotation_matrix(euler_angles):
    phi, theta, psi = euler_angles
    R_x = np.array([[1, 0, 0],
                    [0, np.cos(phi), -np.sin(phi)],
                    [0, np.sin(phi), np.cos(phi)]])

    R_y = np.array([[np.cos(theta), 0, np.sin(theta)],
                    [0, 1, 0],
                    [-np.sin(theta), 0, np.cos(theta)]])

    R_z = np.array([[np.cos(psi), -np.sin(psi), 0],
                    [np.sin(psi), np.cos(psi), 0],
                    [0, 0, 1]])

    return R_z @ R_y @ R_x

def dcm_from_euler(euler_angles):
    """
    计算方向余弦矩阵（DCM）
    :param roll: 滚转角，单位为弧度
    :param pitch: 俯仰角，单位为弧度
    :param yaw: 偏航角，单位为弧度
    :return: 方向余弦矩阵（DCM）
    """
    roll, pitch, yaw=euler_angles
    c_roll = np.cos(roll)
    s_roll = np.sin(roll)
    c_pitch = np.cos(pitch)
    s_pitch = np.sin(pitch)
    c_yaw = np.cos(yaw)
    s_yaw = np.sin(yaw)

    # 构建方向余弦矩阵
    dcm = np.array([
        [c_yaw * c_pitch, c_yaw * s_pitch * s_roll - s_yaw * c_roll, c_yaw * s_pitch * c_roll + s_yaw * s_roll],
        [s_yaw * c_pitch, s_yaw * s_pitch * s_roll + c_yaw * c_roll, s_yaw * s_pitch * c_roll - c_yaw * s_roll],
        [-s_pitch, c_pitch * s_roll, c_pitch * c_roll]
    ])
    return dcm

if __name__ == "__main__":
    run_uav_simulation()
